import os
from typing import List
from langchain_openai import ChatOpenAI
from langchain_community.embeddings import DashScopeEmbeddings
from langchain.vectorstores import FAISS


def load_vectorstore(vectorstore_path: str = "faiss_index"):
    """
    加载已有的向量数据库

    Args:
        vectorstore_path (str): 向量数据库存储路径

    Returns:
        FAISS: 加载的向量数据库实例
    """
    # 从配置文件读取API密钥
    with open("config.txt", "r", encoding="utf-8") as f:
        dashscope_api_key = f.read().strip()
    
    # 初始化嵌入模型
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v2",
        dashscope_api_key=dashscope_api_key
    )
    
    # 加载向量数据库
    if os.path.exists(vectorstore_path):
        vectorstore = FAISS.load_local(vectorstore_path, embeddings, allow_dangerous_deserialization=True)
        print(f"向量数据库已从 {vectorstore_path} 加载")
        return vectorstore
    else:
        raise FileNotFoundError(f"向量数据库路径 {vectorstore_path} 不存在")


def search_documents(vectorstore: FAISS, query: str, k: int = 3):
    """
    在向量数据库中搜索相关文档

    Args:
        vectorstore (FAISS): 向量数据库实例
        query (str): 查询语句
        k (int): 返回结果数量

    Returns:
        List[str]: 相关文档内容列表
    """
    # 执行相似性搜索
    docs = vectorstore.similarity_search(query, k=k)
    # 提取文档内容
    contents = [doc.page_content for doc in docs]
    return contents


def parse_format_requirements(context: str):
    """
    解析格式要求并生成样式映射

    Args:
        context (str): 检索到的文档内容

    Returns:
        dict: 样式格式映射
    """
    # 从配置文件读取API密钥
    with open("config.txt", "r", encoding="utf-8") as f:
        dashscope_api_key = f.read().strip()
    
    # 初始化通义千问模型
    llm = ChatOpenAI(
        model="qwen-plus",
        api_key=dashscope_api_key,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )
    
    # 构建提示词
    prompt = f"""
    请分析格式要求，将其转换为样式格式映射字典。
    输出应为一个Python字典，键是样式名（如 'Heading 1', 'Heading 2' 等），值是包含格式设置的字典。
    格式设置字典可以包含以下键：
    - 'font_name' (str): 字体名称 (如 '黑体', '宋体')
    - 'font_size' (int): 字号 (单位：磅)
    - 'bold' (bool): 是否加粗
    - 'alignment' (str): 对齐方式 ('left', 'center', 'right', 'justify')
    - 'space_before' (float): 段前间距 (单位：磅)
    - 'space_after' (float): 段后间距 (单位：磅)
    - 'line_spacing' (float): 行距 (倍数)
    
    只输出Python字典代码，不要有其他文字。

    格式要求:
    {context}

    样式格式映射字典:
    """
    
    # 调用通义千问生成样式映射
    response = llm.predict(prompt)
    
    try:
        # 安全地评估响应为Python字典
        style_mapping = eval(response)
        return style_mapping
    except:
        print("解析格式要求时出错，使用默认样式映射")
        # 返回默认样式映射
        return {
            'Heading 1': {
                'font_name': '黑体',
                'font_size': 14,
                'bold': True,
                'space_before': 9,
                'space_after': 9,
                'alignment': 'center'
            }
        }


def process_document_with_requirements(question: str, vectorstore: FAISS, input_file_path: str = None):
    """
    根据问题检索相关文档内容，解析格式要求并修改Word文档

    Args:
        question (str): 用户问题
        vectorstore (FAISS): 向量数据库实例
        input_file_path (str, optional): 输入文件路径，默认为Input/example.docx
    """
    # 检索相关文档
    relevant_docs = search_documents(vectorstore, question)
    
    # 合并检索到的内容
    context = "\n\n".join(relevant_docs)
    print(f"根据问题'{question}'检索到的相关内容:")
    print(context)
    
    # 解析格式要求
    print("\n正在解析格式要求...")
    style_mapping = parse_format_requirements(context)
    print("解析得到的样式映射:")
    print(style_mapping)
    
    # 设置文件路径
    if input_file_path is None:
        input_doc_path = os.path.join("Input", "example.docx")
    else:
        input_doc_path = input_file_path
        
    output_doc_path = os.path.join("Output", "modified_example.docx")
    
    # 确保输出目录存在
    os.makedirs("Output", exist_ok=True)
    
    # 检查输入文件是否存在
    if not os.path.exists(input_doc_path):
        print(f"输入文件 {input_doc_path} 不存在，请确保文件存在")
        return False
    
    # 导入文档处理工具
    try:
        import sys
        sys.path.append(os.path.join(os.path.dirname(__file__), 'Unit'))
        from Unit.document_tool import modify_document_format
        
        print(f"\n正在修改文档格式: {input_doc_path}")
        success = modify_document_format(input_doc_path, style_mapping, output_doc_path)
        
        if success:
            print(f"文档格式修改成功，已保存至: {output_doc_path}")
        else:
            print("文档格式修改失败")
        
        return success
    except Exception as e:
        print(f"导入或调用文档处理工具时发生错误: {e}")
        return False


def main():
    """
    主函数：加载向量数据库，根据开题报告的要求修改Word文档
    """
    try:
        # 加载向量数据库
        print("正在加载向量数据库...")
        vectorstore = load_vectorstore("Rag/faiss_index")
        print("向量数据库加载成功！")
        
        # 设置问题
        question = "根据开题报告的要求，修改example.docx的格式"
        print(f"\n用户问题: {question}")
        
        # 处理文档
        print("\n开始处理文档...")
        process_document_with_requirements(question, vectorstore)
        
    except FileNotFoundError as e:
        print(f"错误: {e}")
        print("请先运行 vectorize_article.py 生成向量数据库")
    except Exception as e:
        print(f"发生错误: {e}")


if __name__ == "__main__":
    main()